import pickle
import networkx as nx
import matplotlib.pyplot as plt
import numpy as np
from collections import defaultdict
import warnings
warnings.filterwarnings('ignore')

# Read the pickle file
with open('business-payout.pkl', 'rb') as f:
    G = pickle.load(f)

# Analyze the structure
print("="*50)
print("Graph Data Analysis")
print("="*50)
print(f"Data type: {type(G)}")

if isinstance(G, (nx.DiGraph, nx.MultiDiGraph, nx.Graph, nx.MultiGraph)):
    print(f"Total number of nodes: {G.number_of_nodes()}")
    print(f"Total number of edges: {G.number_of_edges()}")

    # Calculate basic statistics
    if G.number_of_nodes() > 0:
        # Degree statistics
        degrees = dict(G.degree())
        avg_degree = sum(degrees.values()) / len(degrees)
        print(f"Average degree: {avg_degree:.2f}")

        # In-degree and out-degree for directed graphs
        if isinstance(G, (nx.DiGraph, nx.MultiDiGraph)):
            in_degrees = dict(G.in_degree())
            out_degrees = dict(G.out_degree())
            avg_in_degree = sum(in_degrees.values()) / len(in_degrees)
            avg_out_degree = sum(out_degrees.values()) / len(out_degrees)
            print(f"Average in-degree: {avg_in_degree:.2f}")
            print(f"Average out-degree: {avg_out_degree:.2f}")

        # Find most connected nodes
        sorted_nodes = sorted(degrees.items(), key=lambda x: x[1], reverse=True)
        print("\nTop 10 most connected nodes (by total degree):")
        for node, degree in sorted_nodes[:10]:
            if isinstance(G, (nx.DiGraph, nx.MultiDiGraph)):
                print(f"  {node}: degree={degree}, in={in_degrees[node]}, out={out_degrees[node]}")
            else:
                print(f"  {node}: {degree} connections")

        # Graph connectivity
        print("\nGraph Statistics:")
        if isinstance(G, (nx.DiGraph, nx.MultiDiGraph)):
            if nx.is_weakly_connected(G):
                print("  Graph is weakly connected")
            else:
                num_components = nx.number_weakly_connected_components(G)
                print(f"  Number of weakly connected components: {num_components}")

                # Get component sizes
                components = list(nx.weakly_connected_components(G))
                component_sizes = sorted([len(c) for c in components], reverse=True)
                print(f"  Component sizes (top 5): {component_sizes[:5]}")
        else:
            if nx.is_connected(G):
                print("  Graph is connected")
            else:
                num_components = nx.number_connected_components(G)
                print(f"  Number of connected components: {num_components}")

        # Visualize the graph
        # Select nodes for visualization (sample if too large)
        if G.number_of_nodes() > 200:
            # Get the most connected nodes for visualization
            top_nodes = [node for node, _ in sorted_nodes[:100]]
            # Add their neighbors to create a meaningful subgraph
            nodes_to_viz = set(top_nodes)
            for node in top_nodes[:20]:  # Add neighbors of top 20 nodes
                nodes_to_viz.update(G.neighbors(node))
                if isinstance(G, (nx.DiGraph, nx.MultiDiGraph)):
                    nodes_to_viz.update(G.predecessors(node))
            nodes_to_viz = list(nodes_to_viz)[:200]  # Limit to 200 nodes
            subgraph = G.subgraph(nodes_to_viz)
            print(f"\nVisualizing subgraph with {len(nodes_to_viz)} nodes (sampled from {G.number_of_nodes()} total)")
        else:
            subgraph = G
            print(f"\nVisualizing entire graph with {G.number_of_nodes()} nodes")

        # Create main visualization
        plt.figure(figsize=(16, 12))

        # Choose layout based on graph size
        if subgraph.number_of_nodes() < 50:
            pos = nx.spring_layout(subgraph, k=3, iterations=100)
        else:
            pos = nx.spring_layout(subgraph, k=1, iterations=50)

        # Calculate node sizes based on degree
        node_degrees = dict(subgraph.degree())
        max_degree = max(node_degrees.values()) if node_degrees else 1
        node_sizes = [300 + 1000 * (node_degrees[node] / max_degree) for node in subgraph.nodes()]

        # Color nodes by degree
        node_colors = [node_degrees[node] for node in subgraph.nodes()]

        # Draw nodes
        nodes = nx.draw_networkx_nodes(subgraph, pos,
                                       node_color=node_colors,
                                       node_size=node_sizes,
                                       cmap='YlOrRd',
                                       alpha=0.8,
                                       vmin=0,
                                       vmax=max_degree)

        # Draw edges
        nx.draw_networkx_edges(subgraph, pos,
                              edge_color='gray',
                              arrows=True,
                              alpha=0.3,
                              arrowsize=8,
                              width=0.5,
                              connectionstyle='arc3,rad=0.1')

        # Add labels for high-degree nodes
        high_degree_threshold = sorted(node_degrees.values(), reverse=True)[min(15, len(node_degrees)-1)]
        high_degree_nodes = [node for node in subgraph.nodes()
                            if node_degrees[node] >= high_degree_threshold]

        labels = {}
        for node in high_degree_nodes:
            if '/' in str(node):
                labels[node] = str(node).split('/')[-1][:20]
            else:
                labels[node] = str(node)[:20]

        nx.draw_networkx_labels(subgraph, pos, labels,
                               font_size=8,
                               font_weight='bold')

        plt.title(f"Business Payout Dependency Graph\n({subgraph.number_of_nodes()} nodes, {subgraph.number_of_edges()} edges)",
                  fontsize=16)
        plt.colorbar(nodes, label='Node Degree')
        plt.axis('off')
        plt.tight_layout()
        plt.savefig('graph_visualization.png', dpi=300, bbox_inches='tight')
        print("Main visualization saved as 'graph_visualization.png'")
        plt.show()

        # Create degree distribution plot
        plt.figure(figsize=(12, 5))

        # Subplot 1: Degree distribution histogram
        plt.subplot(1, 2, 1)
        all_degrees = [d for n, d in G.degree()]
        plt.hist(all_degrees, bins=min(50, len(set(all_degrees))),
                alpha=0.7, color='blue', edgecolor='black')
        plt.xlabel('Degree', fontsize=12)
        plt.ylabel('Frequency', fontsize=12)
        plt.title('Degree Distribution', fontsize=14)
        plt.grid(True, alpha=0.3)
        plt.yscale('log')

        # Subplot 2: In-degree vs Out-degree scatter (for directed graphs)
        if isinstance(G, (nx.DiGraph, nx.MultiDiGraph)):
            plt.subplot(1, 2, 2)
            in_deg = [G.in_degree(n) for n in G.nodes()]
            out_deg = [G.out_degree(n) for n in G.nodes()]
            plt.scatter(in_deg, out_deg, alpha=0.5, s=20)
            plt.xlabel('In-degree', fontsize=12)
            plt.ylabel('Out-degree', fontsize=12)
            plt.title('In-degree vs Out-degree', fontsize=14)
            plt.grid(True, alpha=0.3)

            # Add diagonal line for reference
            max_val = max(max(in_deg), max(out_deg))
            plt.plot([0, max_val], [0, max_val], 'r--', alpha=0.3)

        plt.tight_layout()
        plt.savefig('degree_analysis.png', dpi=300, bbox_inches='tight')
        print("Degree analysis saved as 'degree_analysis.png'")
        plt.show()

        # Print edge information sample
        print("\nSample edges (first 10):")
        edge_count = 0
        for u, v, data in list(G.edges(data=True))[:10]:
            print(f"  {u} -> {v}")
            if data:
                print(f"    Attributes: {data}")
            edge_count += 1
            if edge_count >= 10:
                break